package grpc

import (
	"context"
	"github.com/ecodeclub/ekit/queue"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"sync"
	"sync/atomic"
	"time"
)

type CounterLimiter struct {
	cnt       atomic.Int32
	threshold int32
}

// 计数器
func (c *CounterLimiter) BuildServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (resp any, err error) {
		c.cnt.Add(1)
		defer func() {
			c.cnt.Add(-1)
		}()
		if c.cnt.Load() <= c.threshold {
			return handler(ctx, req)
		}
		return nil, status.Errorf(codes.ResourceExhausted, "限流")
	}
}

// 固定窗口
type FixWindowLimiter struct {
	//当前窗口的请求数量
	cnt int32
	//窗口允许的最大的请求数量
	threshold int32
	//窗口大小
	window time.Duration
	//上一次窗口的起始时间
	lastStart time.Time
	lock      sync.Mutex
}

func NewFixWindowLimiter(threshold int32, window time.Duration) *FixWindowLimiter {
	return &FixWindowLimiter{cnt: 0, threshold: threshold, window: window, lastStart: time.Now(), lock: sync.Mutex{}}
}

func (f *FixWindowLimiter) BuildServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (resp any, err error) {
		f.lock.Lock()
		now := time.Now()
		//换个窗口
		if now.After(f.lastStart.Add(f.window)) {
			f.lastStart = now
			f.cnt = 1
		}
		f.lock.Unlock()
		if f.cnt <= f.threshold {
			f.cnt++
			// 返回了响应
			return handler(ctx, req)
		}
		return nil, status.Errorf(codes.ResourceExhausted, "限流")
	}
}

// 滑动窗口
type SlidingWindowLimiter struct {
	//窗口允许的最大的请求数量
	threshold int
	//窗口大小
	window time.Duration
	lock   sync.Mutex
	queue  queue.PriorityQueue[time.Time]
}

func NewSlidingWindowLimiter(threshold int, window time.Duration, queue queue.PriorityQueue[time.Time]) *SlidingWindowLimiter {
	return &SlidingWindowLimiter{threshold: threshold, window: window, queue: queue, lock: sync.Mutex{}}
}

func (s *SlidingWindowLimiter) BuildServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
		s.lock.Lock()
		now := time.Now()
		// 快路径检测
		if s.queue.Len() < s.threshold {
			_ = s.queue.Enqueue(now)
			s.lock.Unlock()
			return handler(ctx, handler)
		}

		lastTime := now.Add(-s.window)
		for {
			t, _ := s.queue.Peek()
			if t.Before(lastTime) {
				_, _ = s.queue.Dequeue()
			} else {
				break
			}
		}

		if s.queue.Len() < s.threshold {
			_ = s.queue.Enqueue(now)
			s.lock.Unlock()
			return handler(ctx, handler)
		}

		s.lock.Unlock()
		return nil, status.Errorf(codes.ResourceExhausted, "限流")
	}
}

// 令牌
type TokenBucketLimiter struct {
	//隔多久产生一个令牌
	ticket    time.Ticker
	buckets   chan struct{}
	closeCh   chan struct{}
	closeOnce *sync.Once
}

func NewTokenBucketLimiter(ticket time.Ticker, buckets chan struct{}) *TokenBucketLimiter {
	return &TokenBucketLimiter{ticket: ticket, buckets: buckets, closeCh: make(chan struct{}),
		closeOnce: &sync.Once{}}
}

func (t *TokenBucketLimiter) BuildServerInterceptor() grpc.UnaryServerInterceptor {
	ticker := t.ticket
	go func() {
		for {
			select {
			case <-ticker.C:
				select {
				case t.buckets <- struct{}{}:
				default:
					//bucket满了，忍了
				}
			case <-t.closeCh:
				return
			}
		}
	}()
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (resp any, err error) {
		select {
		case <-t.buckets:
			return handler(ctx, req)
		default:
			return nil, status.Errorf(codes.ResourceExhausted, "限流")
		}
	}
}
func (t *TokenBucketLimiter) close() {
	t.closeOnce.Do(func() {
		close(t.closeCh)
	})
}
